import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import ast
import warnings
from pathlib import Path
import matplotlib.pyplot as plt
import joblib
import sys
sys.path.append(str(Path(__file__).resolve().parent.parent))  # 添加项目根目录
from example.machine_learning.autoencoder import Autoencoder

warnings.filterwarnings('ignore')

# 设置随机种子以确保可重复性
torch.manual_seed(42)
np.random.seed(42)
torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 1. 读取数据
def load_and_preprocess_data(file_path):
    """读取CSV文件并预处理数据"""
    df = pd.read_csv(file_path)
    print(f"原始数据形状: {df.shape}")
    
    # 显示缺失值情况
    print("缺失值统计:")
    print(df.isnull().sum())
    
    # 创建数据副本用于处理
    df_processed = df.copy()
    
    # 处理列表类型的列（RSSI数据）
    list_columns = []
    for col in df.columns:
        if df[col].dtype == object and df[col].notna().any(): #notna表示将该列转化为bool series，非空值为1，any函数检测series中是否有true，有则返回1
            # 检查是否是列表格式的数据
            sample_val = df[col].dropna().iloc[0] #取这一列中第一个非空值
            if isinstance(sample_val, str) and sample_val.startswith('[') and sample_val.endswith(']'):
                list_columns.append(col)
                print(f"处理列表列: {col}")
                
                # 解析列表并提取统计特征   apply函数表示对每个元素应用一个函数（此处是lambda函数,冒号左边是无名函数的参数，右边是函数表达式）
                df_processed[f'{col}_mean'] = df[col].apply(lambda x: np.mean(ast.literal_eval(x)) if pd.notna(x) and x != '' else np.nan) #ast.literal_eval将字符串解析为python对象列表
                df_processed[f'{col}_std'] = df[col].apply(lambda x: np.std(ast.literal_eval(x)) if pd.notna(x) and x != '' else np.nan)
                df_processed[f'{col}_max'] = df[col].apply(lambda x: np.max(ast.literal_eval(x)) if pd.notna(x) and x != '' else np.nan)
                df_processed[f'{col}_min'] = df[col].apply(lambda x: np.min(ast.literal_eval(x)) if pd.notna(x) and x != '' else np.nan)
        if df[col].isna().all():#删除全空列
            print("找到了全空列: ", col)
            df_processed = df_processed.drop(columns=col)
            
    # 删除原始列表列
    df_processed = df_processed.drop(columns=list_columns)
    
    # 选择数值列
    numeric_cols = df_processed.select_dtypes(include=[np.number]).columns.tolist() #select_dtypes是pandas方法，用于按照类型选择列
    df_numeric = df_processed[numeric_cols]
    
    print(f"处理后数据形状: {df_numeric.shape}")
    return df_numeric, df_processed, numeric_cols

def impute_with_trained_autoencoder(model, scaler:StandardScaler, df_original:pd.DataFrame):
    """
    使用训练好的自编码器进行最终填补
    df_original: 原始数据框，包含缺失值
    column_means: 每列的均值，用于初始填充
    """
    df_temp = df_original.copy()
    # 创建临时填充数据（用于标准化）
    for col in df_original.columns:
        if df_original[col].isnull().any():
            df_temp[col] = df_original[col].fillna(df_original[col].mean())
    
    # 标准化
    X_temp = df_temp.values
    X_scaled = scaler.transform(X_temp)
    
    # 使用自编码器进行预测
    model.eval()
    with torch.no_grad():
        X_tensor = torch.FloatTensor(X_scaled).to(model.device)
        X_reconstructed = model(X_tensor).detach().numpy()
    
    # 反标准化
    X_imputed = scaler.inverse_transform(X_reconstructed)
    
    # 创建最终结果：只在原始缺失位置使用预测值
    mask = df_original.isnull().to_numpy()
    print(np.sum(mask))
    print(X_imputed)
    df_final = df_original.to_numpy()
    df_final[mask] = X_imputed[mask] #布尔掩码索引，values返回的是视图，修改的是dataframe的值，但to_numpy()返回的是副本，修改副本不会影响原始dataframe
    print(f"编码器输出{X_imputed[mask]}")
    print(f"填补后{df_final[mask]}")
    return pd.DataFrame(df_final, columns=df_original.columns)

# 主函数
def main():
    # 读取数据
    file_path = Path(__file__).resolve().parent / 'data' / 'training_set_3ap_loc30_nav82.csv'
    df_numeric, df_processed, numeric_cols = load_and_preprocess_data(file_path)
    df_numeric.to_csv(Path(__file__).resolve().parent / 'processed_data.csv', index=False)
    # 创建缺失值掩码
    mask = df_numeric.isnull().to_numpy() #isnull函数返回一个相同形状的布尔型DataFrame，表示每个位置是否为缺失值。.values表示将dataframe转化为numpy数组
    print(f"缺失值总数: {np.sum(mask)}")
    # 标准化数据
    scaler = StandardScaler()
    df_filled_temp = df_numeric.copy()
    # 先用均值填充临时值以便标准化（和取出非缺失数据再标准化是一样的）
    for col in df_numeric.columns:
        if df_numeric[col].isnull().any():
            df_filled_temp[col] = df_numeric[col].fillna(df_numeric[col].mean())
    # 标准化
    X_scaled = scaler.fit_transform(df_filled_temp.to_numpy())
    # 准备训练数据（用填充后的数据训练编码器，然后输出重新补全，只在缺失位置使用预测值）
    X_train, X_val = train_test_split(X_scaled, test_size=0.2, random_state=42)
    train_dataset = TensorDataset(torch.FloatTensor(X_train)) #创建数据库对象
    val_dataset = TensorDataset(torch.FloatTensor(X_val))
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) #创建数据加载器对象
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
    model = Autoencoder(input_dim=X_scaled.shape[1], encoding_dim=16)
    # 训练自编码器
    scalar_path = Path(__file__).resolve().parent / 'model' / 'scaler.joblib'
    model_path = Path(__file__).resolve().parent / 'model' / 'autoencoder.pth'
    joblib.dump(scaler, scalar_path)# 保存标准化器
    train_losses, val_losses = model.train_autoencoder(train_loader, val_loader, num_epochs=100, learning_rate=0.001, save_path = model_path)
    # 画损失图
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    fig, (ax1, ax2) = plt.subplots(nrows = 2,ncols=1,figsize =(8,6))# plt三层管理，首先最上层plt创建总图fig和子图ax1ax2
    ax1.plot(np.arange(0,100), np.array(train_losses))
    ax1.set_title('训练损失')
    ax1.set_xlabel('epoch')
    ax1.set_ylabel('train loss')
    ax1.grid(True)
    ax2.plot(np.arange(0,100), np.array(val_losses))
    ax2.set_title('验证损失')
    ax2.set_xlabel('epoch')
    ax2.set_ylabel('val loss')
    ax2.grid(True)
    plt.savefig(Path(__file__).resolve().parent / 'autoencoder_loss_curve.png')
    # 利用训练好的自编码器进行最终填补
    df_final = impute_with_trained_autoencoder(model, scaler, df_numeric)
    df_final.to_csv(Path(__file__).resolve().parent / 'completed_data.csv', index=False)
    return

# 运行主函数
if __name__ == "__main__":
    df_result = main()